from DB import db, r
import pickle


def transaction(func):
    def inner(*args, **kwargs):
        try:
            res = func(*args, **kwargs)
            return res
        except Exception as e:
            db.session.rollback()
            raise e

    return inner


def cache(forever=False, ex=360, key=None):
    def outer(func):
        cache_key = key if key else func.__name__

        def inner(*args, **kwargs):
            data = r.get(cache_key)
            if data:
                return pickle.loads(data)
            else:
                res = func(*args, **kwargs)
                data = pickle.dumps(res)
                if not forever:
                    r.set(cache_key, data, ex)
                else:
                    r.set(cache_key, data)
                    r.persist(cache_key)
                return res

        return inner

    return outer
